# ---------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# ---------------------------------------------------------

import logging
import os
from typing import Dict, List, Optional, Union

from azure.core.credentials import TokenCredential
from azure.core.pipeline.policies import UserAgentPolicy
from azure.ai.evaluation._model_configurations import AzureAIProject
from azure.ai.evaluation.simulator._model_tools import ManagedIdentityAPITokenManager
from azure.ai.evaluation._common.raiclient import MachineLearningServicesClient
from azure.ai.evaluation._constants import TokenScope
from azure.ai.evaluation._common.utils import is_onedp_project
from azure.ai.evaluation._common.onedp import ProjectsClient as AIProjectClient
from azure.ai.evaluation._common import EvaluationServiceOneDPClient
from azure.ai.evaluation._user_agent import UserAgentSingleton
import jwt
import time
import ast


class GeneratedRAIClient:
    """Client for the Responsible AI Service using the auto-generated MachineLearningServicesClient.

    :param azure_ai_project: The Azure AI project, which can either be a string representing the project endpoint
        or an instance of AzureAIProject. It contains subscription id, resource group, and project name.
    :type azure_ai_project: Union[str, ~azure.ai.evaluation.AzureAIProject]
    :param token_manager: The token manager
    :type token_manager: ~azure.ai.evaluation.simulator._model_tools._identity_manager.APITokenManager
    """

    def __init__(
        self,
        azure_ai_project: Union[AzureAIProject, str],
        token_manager: ManagedIdentityAPITokenManager,
    ):
        self.azure_ai_project = azure_ai_project
        self.token_manager = token_manager
        self.logger = logging.getLogger(__name__)

        user_agent_policy = UserAgentPolicy(base_user_agent=UserAgentSingleton().value)

        if not is_onedp_project(azure_ai_project):
            # Service URL construction
            if "RAI_SVC_URL" in os.environ:
                endpoint = os.environ["RAI_SVC_URL"].rstrip("/")
            else:
                endpoint = self._get_service_discovery_url()

            # Create the autogenerated client
            self._client = MachineLearningServicesClient(
                endpoint=endpoint,
                subscription_id=self.azure_ai_project["subscription_id"],
                resource_group_name=self.azure_ai_project["resource_group_name"],
                workspace_name=self.azure_ai_project["project_name"],
                credential=self.token_manager,
            ).rai_svc
        else:
            self._client = AIProjectClient(
                endpoint=azure_ai_project,
                credential=token_manager,
                user_agent_policy=user_agent_policy,
            ).red_teams
            self._evaluation_onedp_client = EvaluationServiceOneDPClient(
                endpoint=azure_ai_project,
                credential=token_manager,
                user_agent_policy=user_agent_policy,
            )

    def _get_service_discovery_url(self):
        """Get the service discovery URL.

        :return: The service discovery URL
        :rtype: str
        """
        import requests

        bearer_token = self._fetch_or_reuse_token(self.token_manager)
        headers = {
            "Authorization": f"Bearer {bearer_token}",
            "Content-Type": "application/json",
        }

        response = requests.get(
            f"https://management.azure.com/subscriptions/{self.azure_ai_project['subscription_id']}/"
            f"resourceGroups/{self.azure_ai_project['resource_group_name']}/"
            f"providers/Microsoft.MachineLearningServices/workspaces/{self.azure_ai_project['project_name']}?"
            f"api-version=2023-08-01-preview",
            headers=headers,
            timeout=5,
        )

        if response.status_code != 200:
            msg = (
                f"Failed to connect to your Azure AI project. Please check if the project scope is configured "
                f"correctly, and make sure you have the necessary access permissions. "
                f"Status code: {response.status_code}."
            )
            raise Exception(msg)

        # Parse the discovery URL
        from urllib.parse import urlparse

        base_url = urlparse(response.json()["properties"]["discoveryUrl"])
        return f"{base_url.scheme}://{base_url.netloc}"

    async def get_attack_objectives(
        self,
        *,
        risk_type: Optional[str] = None,
        risk_category: Optional[str] = None,
        application_scenario: str = None,
        strategy: Optional[str] = None,
        language: str = "en",
        scan_session_id: Optional[str] = None,
        target: Optional[str] = None,
        client_id: Optional[str] = None,
    ) -> Dict:
        """Get attack objectives using the auto-generated operations.

        :param risk_type: Optional risk type to filter the attack objectives
        :type risk_type: Optional[str]
        :param risk_category: Optional risk category to filter the attack objectives
        :type risk_category: Optional[str]
        :param application_scenario: Optional description of the application scenario for context
        :type application_scenario: str
        :param strategy: Optional strategy to filter the attack objectives
        :type strategy: Optional[str]
        :param language: Language code for the attack objectives (e.g., "en", "es", "fr")
        :type language: str
        :param scan_session_id: Optional unique session ID for the scan
        :type scan_session_id: Optional[str]
        :param target: Optional target type (model/agent)
        :type target: Optional[str]
        :param client_id: Optional client ID for ACA token authorization
        :type client_id: Optional[str]
        :return: The attack objectives
        :rtype: Dict
        """
        try:
            # Build headers dictionary
            headers = {}
            if scan_session_id:
                headers["x-ms-client-request-id"] = scan_session_id
            if client_id:
                from azure.identity import DefaultAzureCredential

                self.logger.info(f"Using client_id: {client_id} to set token in aml-aca-token header ")

                # Get token using the client_id for managed identity
                managed_identity_credential = DefaultAzureCredential(
                    managed_identity_client_id=client_id, exclude_interactive_browser_credential=True
                )
                token = managed_identity_credential.get_token(TokenScope.DEFAULT_AZURE_MANAGEMENT).token
                headers["aml-aca-token"] = token

            # Send the request using the autogenerated client
            response = self._client.get_attack_objectives(
                risk_types=[risk_type],
                risk_category=risk_category,
                lang=language,
                strategy=strategy,
                target_type=target,
                headers=headers,
            )

            return response

        except Exception as e:
            # Log the exception for debugging purposes
            import logging

            logging.error(f"Error in get_attack_objectives: {str(e)}")
            raise

    async def get_jailbreak_prefixes(self, scan_session_id: Optional[str] = None) -> List[str]:
        """Get jailbreak prefixes using the auto-generated operations.

        :param scan_session_id: Optional unique session ID for the scan
        :type scan_session_id: Optional[str]
        :return: The jailbreak prefixes
        :rtype: List[str]
        """
        try:
            # Send the request using the autogenerated client
            response = self._client.get_jail_break_dataset_with_type(
                type="upia", headers={"x-ms-client-request-id": scan_session_id}
            )
            if isinstance(response, list):
                return response
            else:
                self.logger.error("Unexpected response format from get_jail_break_dataset_with_type")
                raise ValueError("Unexpected response format from get_jail_break_dataset_with_type")

        except Exception as e:
            return [""]

    def _fetch_or_reuse_token(self, credential: TokenCredential, token: Optional[str] = None) -> str:
        """Get token. Fetch a new token if the current token is near expiry

        :param credential: The Azure authentication credential.
        :type credential:
        ~azure.core.credentials.TokenCredential
        :param token: The Azure authentication token. Defaults to None. If none, a new token will be fetched.
        :type token: str
        :return: The Azure authentication token.
        """
        if token:
            # Decode the token to get its expiration time
            try:
                decoded_token = jwt.decode(token, options={"verify_signature": False})
            except jwt.PyJWTError:
                pass
            else:
                exp_time = decoded_token["exp"]
                current_time = time.time()

                # Return current token if not near expiry
                if (exp_time - current_time) >= 300:
                    return token

        # Get token
        return credential.get_token(TokenScope.DEFAULT_AZURE_MANAGEMENT).token
